import numpy as np
import random
import matplotlib.pyplot as plt
from ft import ft  # 外包络函数
from designrs import designrs  # 设计谱函数
from resspec import resspec  # 响应谱函数
from hecheng import hecheng  # 地震波合成函数
from find_bad_w import find_bad_w  # 找误差最大的频率
import time
import geatpy as ea  # import geatpy
import time


class MyProblem(ea.Problem):  # 继承Problem父类
    def __init__(self):
        name = 'MyProblem'  # 初始化name（函数名称，可以随意设置）
        M = 1  # 初始化M（目标维数）
        maxormins = [1]  # 初始化maxormins（目标最小最大化标记列表，1：最小化该目标；-1：最大化该目标）
        Dim = 3  # 初始化Dim（决策变量维数）
        varTypes = [0] * Dim  # 初始化varTypes（决策变量的类型，元素为0表示对应的变量是连续的；1表示是离散的）
        lb = [0.2] * Dim  # 决策变量下界
        ub = [2] * Dim  # 决策变量上界
        lbin = [1] * Dim  # 决策变量下边界（0表示不包含该变量的下边界，1表示包含）
        ubin = [1] * Dim  # 决策变量上边界（0表示不包含该变量的上边界，1表示包含）
        # 调用父类构造方法完成实例化
        ea.Problem.__init__(self, name, M, maxormins, Dim, varTypes, lb, ub, lbin, ubin)

    def aimFunc(self, pop):  # 目标函数
        Vars = pop.Phen  # 得到决策变量矩阵
        global t, Aw, fai, bl, w, index_w, Tyz, Damp, sjp
        store = []
        for jj in range(10):  # 和种群规模一样大
            x = Vars[jj, 0:3]
            Aw2 = Aw.copy()
            for ii in range(3):
                Aw2[int(index_w[ii])] *= x[ii]
            at_this = hecheng(t, w, fai, Aw2, bl)  # 生成曲线
            MAcc_this = resspec(t, at_this, Tyz, Damp)  # 生成反应谱
            error_this = abs(MAcc_this - sjp) / sjp  # 计算曲线在各周期点误差
            store.append(max(error_this))
        pop.ObjV = np.array(store).reshape(10, 1)  # 赋值给pop种群对象的ObjV属性  矩阵 每一行记录每一个个体的目标函数值


# 超参数
ld = 6  # 地震烈度 6~9
fz = 1  # 场地分组 1~3
cd = 0  # 场地类型 0~4
global Damp
Damp = 0.1  # 阻尼比
rare = 0  # 0多遇地震 1罕遇地震
westeast = 1  # 0东部 1西部

# 基本参数
# 确定需要控制的反应谱Sa（T）(T=T1,...,TM)的坐标点数M， 反应谱控制容差rc
Dt = 0.02  # 特征周期取值间隔
global Tyz
Tyz = np.arange(0.02, 6.02, Dt)  # 反应谱横坐标，即特征周期
nTyz = len(Tyz)  # 特征周期点个数
rc = 0.2  # 容差
g = 9.8  # 重力加速度
delta_f = 0.05  # 频率间隔 Hz
global w
w = np.arange(0.05 * 2 * np.pi, (32 + delta_f) * 2 * np.pi, delta_f * 2 * np.pi)  # 角频率  %频率变化范围0.05Hz~25Hz
npl = len(w)  # 频率点数
P = 0.15  # 反应超越概率
Td = 20  # 人造地震动持续时间
dt = 0.02  # 时间间隔
global t
t = np.arange(0, Td + dt, dt)  # 人造地震动时间序列
nt = len(t)
global sjp
sjp = np.zeros(nTyz)
for i in range(nTyz):  # 得到一定场地条件下的设计谱
    sjp[i] = g * designrs(2 * np.pi / Tyz[i], Damp, ld, cd, fz, rare)

# 求幅值系数
# 先把单一频率的分量在各个时间点的幅值求出来，再叠加
# 反应谱转换功率谱
SW = np.zeros(npl)
global Aw
Aw = np.zeros(npl)  # 频率幅值系数
for i in range(npl):
    st = g * designrs(w[i], Damp, ld, cd, fz, rare)  # 设计谱
    SW[i] = - Damp / (np.pi * w[i]) * st ** 2 / np.log(-np.pi / (w[i] * Td) * np.log(1 - P))
    Aw[i] = np.sqrt(4 * SW[i] * delta_f * 2 * np.pi)  # 各频率对应的幅值系数

# # 地震波生成
global fai
fai = np.zeros(npl)  # 随机相位序列
for i in range(npl):  # 频率点
    fai[i] = random.random() * 2 * np.pi  # 随机相位
global bl
bl = np.zeros(nt)  # 外包络
for j in range(nt):
    bl[j] = ft(t[j], fz, ld, westeast)  # 很费时间
# t1 = time.time()
at = hecheng(t, w, fai, Aw, bl)
# t2 = time.time()
# print(t2 - t1)
MAcc = resspec(t, at, Tyz, 0.05)  # 生成初始地震波曲线的反应谱
rcrsp = np.zeros(nTyz)
for i in range(nTyz):
    rcrsp[i] = abs(MAcc[i] - sjp[i]) / sjp[i]  # 计算曲线在各周期点误差

epoch = 1  # 迭代次数计数器
max_epoch = 20  # 简单迭代次数上限
rcrsp_record = [max(rcrsp)]  # 记录误差
Aw_record = [Aw]   # 记录简单迭代过程中的系数
this_Aw = Aw.copy()  # 复制一份 幅值系数
while max(rcrsp) > rc:  # 容差大于设计值
    blxs_1 = sjp / MAcc
    blxs_2 = np.ones(npl)  # 每次都要初始化比例修正系数
    # 去头，将频率换算为周期，若周期<设计谱的最小值，则置为最小值
    blxs_2[np.where(w < Tyz[0])] = blxs_1[1]
    # 去尾，将频率换算为周期，若周期>设计谱的最大值，则置为最大值
    blxs_2[np.where(w > Tyz[-1])] = blxs_1[-1]
    for fp in range(npl):  # 频率点
        # 如果一个频率换算得到的周期不在设计谱的坐标点上，则使用线性插值的方式计算出其幅值调整系数
        this_where = list(np.where(2 * np.pi / w[fp] <= Tyz)[0])
        if this_where:
            tp = this_where[0]  # 找到频率对应的周期在反应谱上的位置
            blxs_2[fp] = blxs_1[tp] + (blxs_1[tp + 1] - blxs_1[tp]) * (2 * np.pi / w[fp] - Tyz[tp]) \
                         / (Tyz[tp + 1] - Tyz[tp])  # 线性插值

    this_Aw = this_Aw * blxs_2  # 修正
    at = hecheng(t, w, fai, this_Aw, bl)
    MAcc = resspec(t, at, Tyz, Damp)  # 生成初始地震波曲线的反应谱
    rcrsp = abs(MAcc - sjp) / sjp  # 计算曲线在各周期点误差
    Aw_record.append(this_Aw)
    rcrsp_record.append(max(rcrsp))
    print(max(rcrsp))
    if epoch > max_epoch:
        break
    epoch += 1
# 选出简单迭代里面最好的 也能加速计算
a = rcrsp_record.index(min(rcrsp_record))  # 找到最好的一代
Aw = Aw_record[a]
at = hecheng(t, w, fai, Aw, bl)
MAcc = resspec(t, at, Tyz, Damp)  # 生成初始地震波曲线的反应谱
for i in range(nTyz):
    rcrsp[i] = abs(MAcc[i] - sjp[i]) / sjp[i]  # 计算曲线在各周期点误差
print(max(rcrsp))


# ga计算
global index_w
index_w = find_bad_w(w, rcrsp, Tyz, 3)  # 找出误差最大的频率点 这是3个

"""===============================实例化问题对象==========================="""
problem = MyProblem()  # 生成问题对象
"""=================================种群设置==============================="""
Encoding = 'BG'  # 编码方式
NIND = 10  # 种群规模
Field = ea.crtfld(Encoding, problem.varTypes, problem.ranges, problem.borders)  # 创建区域描述器
population = ea.Population(Encoding, Field, NIND)  # 实例化种群对象（此时种群还没被初始化，仅仅是完成种群对象的实例化）
# """===============================算法参数设置============================="""
myAlgorithm = ea.soea_SEGA_templet(problem, population)  # 实例化一个算法模板对象
myAlgorithm.MAXGEN = 20  # 最大进化代数

[population, obj_trace, var_trace] = myAlgorithm.run()  # 执行算法模板
# population.save() # 把最后一代种群的信息保存到文件中
# # 输出结果
best_gen = np.argmin(problem.maxormins * obj_trace[:, 1])  # 记录最优种群个体是在哪一代
best_ObjV = obj_trace[best_gen, 1]
print('最优的目标函数值为：%s' % (best_ObjV))
print('最优的控制变量值为：')
for i in range(var_trace.shape[1]):
    print(var_trace[best_gen, i])
    Aw[int(index_w[i])] *= var_trace[best_gen, i]

at = hecheng(t, w, fai, Aw, bl)
MAcc = resspec(t, at, Tyz, Damp)  # 生成初始地震波曲线的反应谱
for i in range(nTyz):
    rcrsp[i] = abs(MAcc[i] - sjp[i]) / sjp[i]  # 计算曲线在各周期点误差
print(max(rcrsp))
plt.plot(rcrsp)
plt.show()

print('有效进化代数：%s' % (obj_trace.shape[0]))
print('最优的一代是第 %s 代' % (best_gen + 1))
print('评价次数：%s' % (myAlgorithm.evalsNum))
print('时间已过 %s 秒' % (myAlgorithm.passTime))

